from torchvision.datasets import CIFAR10
from torchvision import transforms
from torch.utils.data import DataLoader
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
import torch

# 加载CIFAR10数据集
dataset = CIFAR10(root="./dataset", train=False, transform=transforms.ToTensor(), download=True)
# 设置batch_size=64的dataloader
dl = DataLoader(dataset, batch_size=64, shuffle=True)

class Tudui(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, padding=0, stride=1)

    def forward(self, input):
        return self.conv1(input)

writer = SummaryWriter(log_dir="./logs")
count = 1
tudui = Tudui()
for data in dl:
    imgs, targets = data
    print(imgs.shape)
    output = tudui(imgs)
    print(output.shape)

    output = torch.reshape(output, (-1, 3, 30, 30))
    writer.add_images("tudui", output, count)
    count += 1

writer.close()